"""dgcnn test"""
import mindspore as ms
import numpy as np
from mind3d.models.blocks.edgeconv import EdgeConv, Getgraphfeature
from mind3d.utils.knn import KNN
from mindspore import Tensor


def test_EdgeConv(k):
    """
    test the function of EdgeConv
    """
    data = Tensor(np.ones((32, 3, 1024)), ms.float32)
    trans =EdgeConv(layers=[6, 64], K=k, dim9=False)
    out = trans(data)
    assert out.shape == (32, 64, 1024), 'output shape not match'
    print("################################")
    print("EdgeConv test passed!")
    print("################################")

def test_Getgraphfeature(k):
    """
    test the function of Getgraphfeature
    """
    data = Tensor(np.ones((32, 3, 1024)), ms.float32)
    trans=Getgraphfeature(data,k=k,idx=None,dim9=False)
    out=trans(data)
    assert out.shape == (32, 6, 1024, 20), 'output shape not match'
    print("################################")
    print("Getgraphfeature test passed!")
    print("################################")


def test_KNN(k):
    """
    test the function of KNN
    """
    data = Tensor(np.ones((32, 3, 1024)), ms.float32)
    out = KNN(data,k=k)
    print(out.shape)
    assert out.shape == (32, 1024, 20), 'output shape not match'
    print("################################")
    print("KNN test passed!")
    print("################################")


if __name__ == "__main__":
    test_EdgeConv(k=20)
    test_KNN(k=20)
    test_Getgraphfeature(k=20)
